SageMaker HyperPod のログインノードに VS Code から Remote SSH してみる

SageMaker HyperPod のログインノードに VS Code から Remote SSH してみる

#SageMaker HyperPod
#VSCode（Visual Studio Code）
#AWS
たかくに

たかくに

facebook logohatena logotwitter logo
Clock Icon2025.01.10

こんにちは！クラウド事業本部コンサルティング部のたかくに（@takakuni_）です。

SageMaker HyperPod の Slurm WorkShop を見ていると、 Tips に VS Code からログインノードへ Remote SSH するシナリオがありました。

https://catalog.workshops.aws/sagemaker-hyperpod/en-US/05-advanced/05-vs-code

非常に便利だなと思ったため、実際に試してみます。

SageMaker クラスターの構築

まずは HyperPod クラスターを構築します。

今回はログインノードさえ用意できてれば OK のため、最小構成で作成します。

https://github.com/takakuni-classmethod/genai-blog/tree/main/sagemaker_hyperpod_101

EC2 の Remote SSH と比べて

EC2 の Remote SSH と比べて SageMaker HyperPod ならではの部分を解説します。

接続方法は EC2 と同じで、 Systems Manager の仕組みを利用して、Remote SSH を行います。

https://dev.classmethod.jp/articles/how-to-use-vscode-remote-ssh-with-aws-systems-manager/

ただし SageMaker HyperPod の場合、ノードにはキーペアを設定する項目がありません。

そのため、ローカルでキーペアを作成し、サーバーに公開鍵を登録する必要があります。

キーペアの作成

ということで、まずは、ローカルでキーペアを作成します。今回は hyperpod というファイル名にしました。

ssh-keygen -t rsa -q -f "$HOME/.ssh/hyperpod" -N ""

作成した公開鍵を確認します。

うまく作成できていますね。後ほど利用するため、メモしておきましょう。

takakuni@ sagemaker_hyperpod_101 % cat ~/.ssh/hyperpod.pub
ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCwTzpeozioaRNzfjml/iyHxfcYsi+u41LO4VZYVv07HJMf/HTRd3QxMxTu7bES1leQUw3zEItjFTrmqtZZnq292z/ZY/JJbDKv1GomBdd5sRw0POVkOuMOmSE0+3xwlxFG/UXlbUUdrtzLmMUpKCWz2uSOHaHQ0IZ2PGGYYgGTxzjv63xVvMV7aAfAEzU3MtWBUsXF0KKk9ZC0VUiLSIEm6BE3hMKQyUSOeE26+KOiS77KeOaXRlnWwrmsS1GBCPUIqIAz4QEiU82plJsInx1l3jnChgEUpeWBg9I+sLblm6SF4EG/EmT6SnuQt2hQ/XLjFgsTzUJYuMu4L5rU9738TjDn/kMRLeojpgOCksdV4d7TEUTFkeqfOB/xu8sfsfuGC+ozAnITrbQ5d9WC8ZBdDCy7AGxR60kEZqE5YlrBcAtE8iz9xyGSQu3uQqmMZAleH+F1CrizihWpSfg11mgTRfOtNbOdJube/G7bshnXDre2G1mXT4Lw1ysua6B3THM

続いて、サーバーに公開鍵を登録するため、Easy SSH を利用してノードにログインします。

適宜クラスター名やインスタンスグループ名を変更します。

curl -O https://raw.githubusercontent.com/aws-samples/awsome-distributed-training/main/1.architectures/5.sagemaker-hyperpod/easy-ssh.sh
chmod +x easy-ssh.sh
./easy-ssh.sh -c controller-machine sagemaker-hyperpod

https://docs.aws.amazon.com/sagemaker/latest/dg/sagemaker-hyperpod-run-jobs-access-nodes.html#sagemaker-hyperpod-run-jobs-access-nodes-tips

こちらも、うまくログインできていますね。

takakuni@ sagemaker_hyperpod_101 % ./easy-ssh.sh -c controller-machine sagemaker-hyperpod

=================================================

==== 🚀 HyperPod Cluster Easy SSH Script! 🚀 ====


=================================================
Cluster id: cf2jdza8s74w
Instance id: i-05c4ceea296d9faa4
Node Group: controller-machine
grep: /Users/takakuni/.ssh/config: No such file or directory
Would you like to add sagemaker-hyperpod to  ~/.ssh/config (yes/no)?
> no
❌ skipping adding ml-cluster to  ~/.ssh/config:
cat: /Users/takakuni.shinnosuke/.ssh/id_rsa.pub: No such file or directory
1. Detected SSH public key ~/.ssh/id_rsa.pub on the cluster. Skipping adding...

Now you can run:

$ ssh sagemaker-hyperpod

Starting session with SessionId: botocore-session-1736514901-ocrg4cseqnydr9io98qiu2pk88
#

Root からの直 SSH 接続は認められていないため、ユーザーを ubuntu に変更します。

# whoami
sudo su - ubuntu
whoamiroot
# ubuntu@ip-172-22-151-244:~$ whoami
ubuntu
ubuntu@ip-172-22-151-244:~$

最後に ~/.ssh/authorized_keys へ、先ほどの公開鍵を登録します。

このセットアップが面倒な場合は、ライフサイクルスクリプトの中で初期設定しても良いかもですね。

~/.ssh/authorized_keys
ssh-rsa AAAAB3NzaC1yc2EAAAADAQABAAABgQCwTzpeozioaRNzfjml/iyHxfcYsi+u41LO4VZYVv07HJMf/HTRd3QxMxTu7bES1leQUw3zEItjFTrmqtZZnq292z/ZY/JJbDKv1GomBdd5sRw0POVkOuMOmSE0+3xwlxFG/UXlbUUdrtzLmMUpKCWz2uSOHaHQ0IZ2PGGYYgGTxzjv63xVvMV7aAfAEzU3MtWBUsXF0KKk9ZC0VUiLSIEm6BE3hMKQyUSOeE26+KOiS77KeOaXRlnWwrmsS1GBCPUIqIAz4QEiU82plJsInx1l3jnChgEUpeWBg9I+sLblm6SF4EG/EmT6SnuQt2hQ/XLjFgsTzUJYuMu4L5rU9738TjDn/kMRLeojpgOCksdV4d7TEUTFkeqfOB/xu8sfsfuGC+ozAnITrbQ5d9WC8ZBdDCy7AGxR60kEZqE5YlrBcAtE8iz9xyGSQu3uQqmMZAleH+F1CrizihWpSfg11mgTRfOtNbOdJube/G7bshnXDre2G1mXT4Lw1ysua6B3THM=

セットアップ完了です。 exit で抜けちゃいます。

ubuntu@ip-172-22-151-244:~$ exit
logout
# exit

Exiting session with sessionId: botocore-session-1736514901-ocrg4cseqnydr9io98qiu2pk88.

takakuni@ sagemaker_hyperpod_101 %

ローカルの ~/.ssh/config を編集します。クラスター名やリージョン等は適宜変更します。

cat <<EOF >> ~/.ssh/config
Host sagemaker-hyperpod
  Hostname sagemaker-cluster:cf2jdza8s74w_controller-machine-i-05c4ceea296d9faa4
  User ubuntu
  Port 22
  IdentityFile ~/.ssh/hyperpod
  ProxyCommand sh -c "aws ssm start-session --target %h --region ap-northeast-1 --document-name AWS-StartSSHSession --parameters 'portNumber=%p' --profile takakuni-tf"
EOF

ローカルから ~/.ssh/config を利用して、接続テストを行います。

ssh sagemaker-hyperpod

うまく設定できていれば、以下のように接続できます。

takakuni@ sagemaker_hyperpod_101 % ssh sagemaker-hyperpod
========================================================================================
    ____              __  ___     __             __ __                  ___          __
   / __/__ ____ ____ /  |/  /__ _/ /_____ ____  / // /_ _____  ___ ____/ _ \___  ___/ /
  _\ \/ _ `/ _ `/ -_) /|_/ / _ `/  '_/ -_) __/ / _  / // / _ \/ -_) __/ ___/ _ \/ _  /
 /___/\_,_/\_, /\__/_/  /_/\_,_/_/\_\\__/_/   /_//_/\_, / .__/\__/_/ /_/   \___/\_,_/
          /___/                                    /___/_/
                          HyperPod Instance AMI (Ubuntu 20.04)
========================================================================================

Welcome to Ubuntu 20.04.6 LTS (GNU/Linux 5.15.0-1072-aws x86_64v)


Utility libraries are installed in /usr/bin/python3.9.
To access them, use /usr/bin/python3.9.

AWS Deep Learning AMI Homepage: https://aws.amazon.com/machine-learning/amis/
Release Notes: https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html
Support: https://forums.aws.amazon.com/forum.jspa?forumID=263
For a fully managed experience, check out Amazon SageMaker at https://aws.amazon.com/sagemaker
=============================================================================

 * Documentation:  https://help.ubuntu.com
 * Management:     https://landscape.canonical.com
 * Support:        https://ubuntu.com/pro

 System information as of Fri Jan 10 14:15:00 UTC 2025

  System load:  0.09               Processes:             192
  Usage of /:   52.3% of 96.73GB   Users logged in:       0
  Memory usage: 28%                IPv4 address for ens6: 172.22.151.244
  Swap usage:   0%

 * Ubuntu Pro delivers the most comprehensive open source security and
   compliance features.

   https://ubuntu.com/aws/pro

Expanded Security Maintenance for Applications is not enabled.

19 updates can be applied immediately.
11 of these updates are standard security updates.
To see these additional updates run: apt list --upgradable

41 additional security updates can be applied with ESM Apps.
Learn more about enabling ESM Apps service at https://ubuntu.com/esm


To replace an instance run:
   sudo scontrol update node=<hostname> state=fail reason="Action:Replace"

To automatically resume jobs, please add the following in your job submission script:
   srun --auto-resume=1

Instance Type: t3.medium

=============================================================================
AMI Name: Deep Learning Base OSS Nvidia Driver GPU AMI (Ubuntu 20.04)
Supported EC2 instances: G4dn, G5, G6, Gr6, G6e, P4d, P4de, P5, P5e, P5en, Trn1, Trn1n
NVIDIA driver version: 550.127.05
CUDA versions available: cuda-12.1 cuda-12.2 cuda-12.3 cuda-12.4
Default CUDA version is 12.1

Release notes: https://docs.aws.amazon.com/dlami/latest/devguide/appendix-ami-release-notes.html
AWS Deep Learning AMI Homepage: https://aws.amazon.com/machine-learning/amis/
Developer Guide and Release Notes: https://docs.aws.amazon.com/dlami/latest/devguide/what-is-dlami.html
Support: https://forums.aws.amazon.com/forum.jspa?forumID=263
For a fully managed experience, check out Amazon SageMaker at https://aws.amazon.com/sagemaker
=============================================================================

The programs included with the Ubuntu system are free software;
the exact distribution terms for each program are described in the
individual files in /usr/share/doc/*/copyright.

Ubuntu comes with ABSOLUTELY NO WARRANTY, to the extent permitted by
applicable law.

   ____              __  ___     __             __ __                  ___          __
  / __/__ ____ ____ /  |/  /__ _/ /_____ ____  / // /_ _____  ___ ____/ _ \___  ___/ /
 _\ \/ _ `/ _ `/ -_) /|_/ / _ `/  '_/ -_) __/ / _  / // / _ \/ -_) __/ ___/ _ \/ _  /
/___/\_,_/\_, /\__/_/  /_/\_,_/_/\_\\__/_/   /_//_/\_, / .__/\__/_/ /_/   \___/\_,_/
         /___/                                    /___/_/

To replace an instance run:
   sudo scontrol update node=<hostname> state=fail reason="Action:Replace"

To automatically resume jobs, please add the following in your job submission script:
   srun --auto-resume=1

You're on the controller
Instance Type: ml.t3.medium
ubuntu@ip-172-22-151-244:~$

それでは最後に、肝心の Remote SSH を実行します。

VS Code 左下の Remote SSH をクリックします。

2025-01-10 at 23.25.17-.envrc — genai-blog@2x.png

ホストに接続するから、先ほど構成したホストを選択しましょう。

2025-01-10 at 23.18.33-.envrc — genai-blog@2x.png

ここでうまく接続いかない場合は、プロファイルを付けてあげるなどしてみてください。

~/.ssh/config
Host sagemaker-hyperpod
  Hostname sagemaker-cluster:cf2jdza8s74w_controller-machine-i-05c4ceea296d9faa4
  User ubuntu
  Port 22
  IdentityFile ~/.ssh/hyperpod
+ ProxyCommand sh -c "aws ssm start-session --target %h --region ap-northeast-1 --document-name AWS-StartSSHSession --parameters 'portNumber=%p' --profile takakuni"
- ProxyCommand sh -c "aws ssm start-session --target %h --region ap-northeast-1 --document-name AWS-StartSSHSession --parameters 'portNumber=%p'"

Remote SSH で SSH 接続できました。

サーバー内部のファイルも認識してローカルから VS Code を介して触れるようになっています。非常に便利ですね。

2025-01-10 at 23.22.41-authorized_keys — ubuntu [SSH sagemaker-hyperpod]@2x.png

まとめ

以上、「SageMaker HyperPod のログインノードに Remote SSH してみた」でした。

普段の EC2 の Remote SSH 開通に比べて、公開鍵の登録がワンステップ必要ですが、それ以外は同じような操作感で繋げられることがわかりました。

非常に便利な機能ですので、ぜひ活用していきたいです。このブログがどなたかの参考になれば幸いです。

クラウド事業本部コンサルティング部のたかくに（@takakuni_）でした！

Share this article

facebook logohatena logotwitter logo

関連記事

Amazon CloudWatch Observability アドオンを利用して SageMaker HyperPod を監視してみた

Amazon CloudWatch Observability アドオンを利用して SageMaker HyperPod を監視してみた

User avatar
たかくに
2025.01.07
EKS オーケストレータを使った SageMaker HyperPod クラスターで S3 をマウントしてみた

EKS オーケストレータを使った SageMaker HyperPod クラスターで S3 をマウントしてみた

User avatar
たかくに
2025.01.06
EKS オーケストレータを使った SageMaker HyperPod クラスターで FSx for Lustre をマウントしてみた

EKS オーケストレータを使った SageMaker HyperPod クラスターで FSx for Lustre をマウントしてみた

User avatar
たかくに
2025.01.05
SageMaker HyperPod の 自動ノード復旧を試してみる

SageMaker HyperPod の 自動ノード復旧を試してみる

User avatar
たかくに
2025.01.04
Classmethod's white logo
Facebook's logo
X's logo
YouTube's logo

© Classmethod, Inc. All rights reserved.